from fastai.vision.all import L


def dfs_change_act(parent_module, name, child_module, target_type, sub_type, _count):
    if child_module is None:
        return
    if type(child_module) == target_type:
        setattr(parent_module, name, sub_type())
        _count[0] += 1
    else:
        for name, module in child_module.named_children():
            dfs_change_act(child_module, name, module, target_type, sub_type, _count)


def change_act(learner, target_type, sub_type):
    count = [0]
    _model = learner.model
    for name, module in _model.named_children():
        dfs_change_act(_model, name, module, target_type, sub_type, count)
    learner._logger.info(f'Change {count} activation unit from {str(target_type)} to {str(sub_type)}')


def multi_head_model_splitter(model):
    feature = L([p for p in model.feature.parameters()])
    head = L([p for p in model.head1.parameters()]) + L([p for p in model.head2.parameters()])
    return L(feature, head)


def simple_head_splitter(model):
    head = [p for p in model.fc.parameters()]
    head_id = [id(p) for p in head]
    feature = [p for p in model.parameters() if id(p) not in head_id]
    return L(L(feature), L(head))
